import pika
import logging
import time
from typing import Dict, Any, Callable, Optional
from pika.exceptions import AMQPConnectionError, ChannelClosedByBroker

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('rabbitmq_client')

class RabbitMQProducer:
    """RabbitMQ 消息生产者"""
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.connection = None
        self.channel = None
        self._connect()
        
    def _connect(self) -> None:
        """建立到 RabbitMQ 的连接"""
        try:
            credentials = pika.PlainCredentials(
                self.config['RABBITMQ_USERNAME'], 
                self.config['RABBITMQ_PASSWORD']
            )
            
            params = pika.ConnectionParameters(
                host=self.config['RABBITMQ_HOST'],
                port=self.config['RABBITMQ_PORT'],
                virtual_host=self.config['RABBITMQ_VHOST'],
                credentials=credentials,
                connection_attempts=self.config['PRODUCER']['connection_attempts'],
                retry_delay=self.config['PRODUCER']['retry_delay']
            )
            
            self.connection = pika.BlockingConnection(params)
            self.channel = self.connection.channel()
            
            # 声明交换器
            for exchange_name, exchange_opts in self.config['EXCHANGES'].items():
                self.channel.exchange_declare(
                    exchange=exchange_name,
                    exchange_type=exchange_opts['type'],
                    durable=exchange_opts.get('durable', True)
                )
            
            # 声明队列
            for queue_name, queue_opts in self.config['QUEUES'].items():
                self.channel.queue_declare(
                    queue=queue_name,
                    durable=queue_opts.get('durable', True),
                    arguments={
                        'x-max-priority': queue_opts.get('max_priority')
                    } if queue_opts.get('max_priority') else None
                )
            
            # 绑定队列到交换器
            for queue_name, binding_opts in self.config['BINDINGS'].items():
                self.channel.queue_bind(
                    queue=queue_name,
                    exchange=binding_opts['exchange'],
                    routing_key=binding_opts['routing_key']
                )
            
            # 启用发布确认
            if self.config['PRODUCER']['confirm_delivery']:
                self.channel.confirm_delivery()
                
            logger.info("成功连接到 RabbitMQ 服务器")
            
        except AMQPConnectionError as e:
            logger.error(f"连接 RabbitMQ 失败: {e}")
            raise
    
    def publish(self, exchange: str, routing_key: str, message: str, 
                properties: Optional[Dict[str, Any]] = None, 
                mandatory: bool = True) -> bool:
        """
        发布消息到 RabbitMQ
        
        参数:
            exchange: 交换器名称
            routing_key: 路由键
            message: 消息内容
            properties: 消息属性
            mandatory: 是否强制（消息无法路由时返回）
            
        返回:
            发布成功返回 True，失败返回 False
        """
        if not self.connection or self.connection.is_closed:
            self._connect()
        
        try:
            # 构建消息属性
            props = pika.BasicProperties(
                delivery_mode=2,  # 消息持久化
                **(properties or {})
            )
            
            # 发布消息
            self.channel.basic_publish(
                exchange=exchange,
                routing_key=routing_key,
                body=message,
                properties=props,
                mandatory=mandatory
            )
            
            logger.debug(f"消息已发布到 {exchange} [{routing_key}]: {message[:100]}...")
            return True
            
        except (AMQPConnectionError, ChannelClosedByBroker) as e:
            logger.error(f"发布消息失败: {e}")
            self._connect()  # 尝试重新连接
            return False
        except Exception as e:
            logger.error(f"发布消息时发生未知错误: {e}")
            return False
    
    def close(self) -> None:
        """关闭连接"""
        if self.channel and not self.channel.is_closed:
            self.channel.close()
        if self.connection and not self.connection.is_closed:
            self.connection.close()
            logger.info("RabbitMQ 连接已关闭")

class RabbitMQConsumer:
    """RabbitMQ 消息消费者"""
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.connection = None
        self.channel = None
        self.consumer_tag = None
        self._connect()
        
    def _connect(self) -> None:
        """建立到 RabbitMQ 的连接"""
        try:
            credentials = pika.PlainCredentials(
                self.config['RABBITMQ_USERNAME'], 
                self.config['RABBITMQ_PASSWORD']
            )
            
            params = pika.ConnectionParameters(
                host=self.config['RABBITMQ_HOST'],
                port=self.config['RABBITMQ_PORT'],
                virtual_host=self.config['RABBITMQ_VHOST'],
                credentials=credentials
            )
            
            self.connection = pika.BlockingConnection(params)
            self.channel = self.connection.channel()
            
            # 设置 QoS
            self.channel.basic_qos(
                prefetch_count=self.config['CONSUMER']['prefetch_count']
            )
            
            logger.info("成功连接到 RabbitMQ 服务器")
            
        except AMQPConnectionError as e:
            logger.error(f"连接 RabbitMQ 失败: {e}")
            raise
    
    def consume(self, queue: str, callback: Callable, auto_ack: bool = False) -> None:
        """
        开始消费队列中的消息
        
        参数:
            queue: 队列名称
            callback: 消息处理回调函数
            auto_ack: 是否自动确认消息
        """
        if not self.connection or self.connection.is_closed:
            self._connect()
        
        self.consumer_tag = self.channel.basic_consume(
            queue=queue,
            on_message_callback=callback,
            auto_ack=auto_ack
        )
        
        logger.info(f"开始消费队列 '{queue}'，消费者标签: {self.consumer_tag}")
        
        try:
            self.channel.start_consuming()
        except (AMQPConnectionError, ChannelClosedByBroker) as e:
            logger.error(f"消费消息时连接丢失: {e}")
            self._reconnect_and_recover(queue, callback, auto_ack)
        except KeyboardInterrupt:
            logger.info("用户中断，停止消费")
            self.stop_consuming()
    
    def _reconnect_and_recover(self, queue: str, callback: Callable, auto_ack: bool) -> None:
        """重新连接并恢复消费"""
        reconnect_attempts = 0
        max_attempts = self.config['CONSUMER']['max_reconnect_attempts']
        delay = self.config['CONSUMER']['reconnect_delay']
        
        while reconnect_attempts < max_attempts:
            try:
                logger.info(f"尝试重新连接 ({reconnect_attempts + 1}/{max_attempts})")
                self._connect()
                self.consume(queue, callback, auto_ack)
                break
            except AMQPConnectionError:
                reconnect_attempts += 1
                logger.error(f"重新连接失败，{delay} 秒后重试")
                time.sleep(delay)
        else:
            logger.error("达到最大重试次数，停止尝试重新连接")
    
    def stop_consuming(self) -> None:
        """停止消费消息"""
        if self.channel and not self.channel.is_closed and self.consumer_tag:
            self.channel.basic_cancel(self.consumer_tag)
            logger.info("已停止消费消息")
    
    def close(self) -> None:
        """关闭连接"""
        self.stop_consuming()
        
        if self.channel and not self.channel.is_closed:
            self.channel.close()
        if self.connection and not self.connection.is_closed:
            self.connection.close()
            logger.info("RabbitMQ 连接已关闭")    